package com.dyava.commons.token_bucket.redis;

import com.dyava.commons.token_bucket.TokenBucketDefined;
import com.dyava.commons.token_bucket.TokenBucketException;
import com.dyava.commons.token_bucket.TokenBucketPool;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.concurrent.ConcurrentHashMap;

public class RedisTokenBucketPool extends TokenBucketPool {

    ConcurrentHashMap<String,RedisTokenBucket> bucketMap = new ConcurrentHashMap<>();
    StringRedisTemplate stringRedisTemplate;

    public StringRedisTemplate getStringRedisTemplate() {
        return stringRedisTemplate;
    }

    public void setStringRedisTemplate(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    @Override
    public void addDefined(TokenBucketDefined tokenBucketDefined){
        RedisTokenBucket tokenBucket = bucketMap.get(tokenBucketDefined.getName());
        if(tokenBucket==null){
            tokenBucket = new RedisTokenBucket(tokenBucketDefined.getName(),tokenBucketDefined.getLimitNum(),tokenBucketDefined.getLimitDen());
            tokenBucket.setStringRedisTemplate(stringRedisTemplate);
            RedisTokenBucket old = bucketMap.putIfAbsent(tokenBucketDefined.getName(),tokenBucket);
            if(old!=null){
                addDefined(tokenBucketDefined);
            }
        }
    }

    @Override
    public RedisTokenBucket getBucket(String name) {
        RedisTokenBucket tokenBucket = bucketMap.get(name);
        return tokenBucket;
    }

    @Override
    public boolean limit(String name, Object id) {
        RedisTokenBucket tokenBucket = bucketMap.get(name);
        if(tokenBucket==null){
            throw new TokenBucketException(name+"令牌桶不存在");
        }
        return tokenBucket.limit(id);
    }
}
